{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6.5 PyTorchでA2C（Advanced Actor-Critic）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# パッケージのimport\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import gym\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定数の設定\n",
    "ENV = 'CartPole-v0'  # 使用する課題名\n",
    "GAMMA = 0.99  # 時間割引率\n",
    "MAX_STEPS = 200  # 1試行のstep数\n",
    "NUM_EPISODES = 1000  # 最大試行回数\n",
    "\n",
    "NUM_PROCESSES = 32  # 同時に実行する環境\n",
    "NUM_ADVANCED_STEP = 5  # 何ステップ進めて報酬和を計算するのか設定\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A2Cの損失関数の計算のための定数設定\n",
    "value_loss_coef = 0.5\n",
    "entropy_coef = 0.01\n",
    "max_grad_norm = 0.5\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# メモリクラスの定義\n",
    "\n",
    "\n",
    "class RolloutStorage(object):\n",
    "    '''Advantage学習するためのメモリクラスです'''\n",
    "\n",
    "    def __init__(self, num_steps, num_processes, obs_shape):\n",
    "\n",
    "        self.observations = torch.zeros(num_steps + 1, num_processes, 4)\n",
    "        self.masks = torch.ones(num_steps + 1, num_processes, 1)\n",
    "        self.rewards = torch.zeros(num_steps, num_processes, 1)\n",
    "        self.actions = torch.zeros(num_steps, num_processes, 1).long()\n",
    "\n",
    "        # 割引報酬和を格納\n",
    "        self.returns = torch.zeros(num_steps + 1, num_processes, 1)\n",
    "        self.index = 0  # insertするインデックス\n",
    "\n",
    "    def insert(self, current_obs, action, reward, mask):\n",
    "        '''次のindexにtransitionを格納する'''\n",
    "        self.observations[self.index + 1].copy_(current_obs)\n",
    "        self.masks[self.index + 1].copy_(mask)\n",
    "        self.rewards[self.index].copy_(reward)\n",
    "        self.actions[self.index].copy_(action)\n",
    "\n",
    "        self.index = (self.index + 1) % NUM_ADVANCED_STEP  # インデックスの更新\n",
    "\n",
    "    def after_update(self):\n",
    "        '''Advantageするstep数が完了したら、最新のものをindex0に格納'''\n",
    "        self.observations[0].copy_(self.observations[-1])\n",
    "        self.masks[0].copy_(self.masks[-1])\n",
    "\n",
    "    def compute_returns(self, next_value):\n",
    "        '''Advantageするステップ中の各ステップの割引報酬和を計算する'''\n",
    "\n",
    "        # 注意：5step目から逆向きに計算しています\n",
    "        # 注意：5step目はAdvantage1となる。4ステップ目はAdvantage2となる。・・・\n",
    "        self.returns[-1] = next_value\n",
    "        for ad_step in reversed(range(self.rewards.size(0))):\n",
    "            self.returns[ad_step] = self.returns[ad_step + 1] * \\\n",
    "                GAMMA * self.masks[ad_step + 1] + self.rewards[ad_step]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A2Cのディープ・ニューラルネットワークの構築\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "class Net(nn.Module):\n",
    "\n",
    "    def __init__(self, n_in, n_mid, n_out):\n",
    "        super(Net, self).__init__()\n",
    "        self.fc1 = nn.Linear(n_in, n_mid)\n",
    "        self.fc2 = nn.Linear(n_mid, n_mid)\n",
    "        self.actor = nn.Linear(n_mid, n_out)  # 行動を決めるので出力は行動の種類数\n",
    "        self.critic = nn.Linear(n_mid, 1)  # 状態価値なので出力は1つ\n",
    "\n",
    "    def forward(self, x):\n",
    "        '''ネットワークのフォワード計算を定義します'''\n",
    "        h1 = F.relu(self.fc1(x))\n",
    "        h2 = F.relu(self.fc2(h1))\n",
    "        critic_output = self.critic(h2)  # 状態価値の計算\n",
    "        actor_output = self.actor(h2)  # 行動の計算\n",
    "\n",
    "        return critic_output, actor_output\n",
    "\n",
    "    def act(self, x):\n",
    "        '''状態xから行動を確率的に求めます'''\n",
    "        value, actor_output = self(x)\n",
    "        # dim=1で行動の種類方向にsoftmaxを計算\n",
    "        action_probs = F.softmax(actor_output, dim=1)\n",
    "        action = action_probs.multinomial(num_samples=1)  # dim=1で行動の種類方向に確率計算\n",
    "        return action\n",
    "\n",
    "    def get_value(self, x):\n",
    "        '''状態xから状態価値を求めます'''\n",
    "        value, actor_output = self(x)\n",
    "\n",
    "        return value\n",
    "\n",
    "    def evaluate_actions(self, x, actions):\n",
    "        '''状態xから状態価値、実際の行動actionsのlog確率とエントロピーを求めます'''\n",
    "        value, actor_output = self(x)\n",
    "\n",
    "        log_probs = F.log_softmax(actor_output, dim=1)  # dim=1で行動の種類方向に計算\n",
    "        action_log_probs = log_probs.gather(1, actions)  # 実際の行動のlog_probsを求める\n",
    "\n",
    "        probs = F.softmax(actor_output, dim=1)  # dim=1で行動の種類方向に計算\n",
    "        entropy = -(log_probs * probs).sum(-1).mean()\n",
    "\n",
    "        return value, action_log_probs, entropy\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# エージェントが持つ頭脳となるクラスを定義、全エージェントで共有する\n",
    "import torch\n",
    "from torch import optim\n",
    "\n",
    "\n",
    "class Brain(object):\n",
    "    def __init__(self, actor_critic):\n",
    "        self.actor_critic = actor_critic  # actor_criticはクラスNetのディープ・ニューラルネットワーク\n",
    "        self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=0.01)\n",
    "\n",
    "    def update(self, rollouts):\n",
    "        '''Advantageで計算した5つのstepの全てを使って更新します'''\n",
    "        obs_shape = rollouts.observations.size()[2:]  # torch.Size([4, 84, 84])\n",
    "        num_steps = NUM_ADVANCED_STEP\n",
    "        num_processes = NUM_PROCESSES\n",
    "\n",
    "        values, action_log_probs, entropy = self.actor_critic.evaluate_actions(\n",
    "            rollouts.observations[:-1].view(-1, 4),\n",
    "            rollouts.actions.view(-1, 1))\n",
    "\n",
    "        # 注意：各変数のサイズ\n",
    "        # rollouts.observations[:-1].view(-1, 4) torch.Size([80, 4])\n",
    "        # rollouts.actions.view(-1, 1) torch.Size([80, 1])\n",
    "        # values torch.Size([80, 1])\n",
    "        # action_log_probs torch.Size([80, 1])\n",
    "        # entropy torch.Size([])\n",
    "\n",
    "        values = values.view(num_steps, num_processes,\n",
    "                             1)  # torch.Size([5, 16, 1])\n",
    "        action_log_probs = action_log_probs.view(num_steps, num_processes, 1)\n",
    "\n",
    "        # advantage（行動価値-状態価値）の計算\n",
    "        advantages = rollouts.returns[:-1] - values  # torch.Size([5, 16, 1])\n",
    "\n",
    "        # Criticのlossを計算\n",
    "        value_loss = advantages.pow(2).mean()\n",
    "\n",
    "        # Actorのgainを計算、あとでマイナスをかけてlossにする\n",
    "        action_gain = (action_log_probs*advantages.detach()).mean()\n",
    "        # detachしてadvantagesを定数として扱う\n",
    "\n",
    "        # 誤差関数の総和\n",
    "        total_loss = (value_loss * value_loss_coef -\n",
    "                      action_gain - entropy * entropy_coef)\n",
    "\n",
    "        # 結合パラメータを更新\n",
    "        self.actor_critic.train()  # 訓練モードに\n",
    "        self.optimizer.zero_grad()  # 勾配をリセット\n",
    "        total_loss.backward()  # バックプロパゲーションを計算\n",
    "        nn.utils.clip_grad_norm_(self.actor_critic.parameters(), max_grad_norm)\n",
    "        #  一気に結合パラメータが変化しすぎないように、勾配の大きさは最大0.5までにする\n",
    "\n",
    "        self.optimizer.step()  # 結合パラメータを更新\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 今回エージェントクラスはなしです"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 実行する環境のクラスです\n",
    "import copy\n",
    "\n",
    "\n",
    "class Environment:\n",
    "    def run(self):\n",
    "        '''メインの実行'''\n",
    "\n",
    "        # 同時実行する環境数分、envを生成\n",
    "        envs = [gym.make(ENV) for i in range(NUM_PROCESSES)]\n",
    "\n",
    "        # 全エージェントが共有して持つ頭脳Brainを生成\n",
    "        n_in = envs[0].observation_space.shape[0]  # 状態は4\n",
    "        n_out = envs[0].action_space.n  # 行動は2\n",
    "        n_mid = 32\n",
    "        actor_critic = Net(n_in, n_mid, n_out)  # ディープ・ニューラルネットワークの生成\n",
    "        global_brain = Brain(actor_critic)\n",
    "\n",
    "        # 格納用変数の生成\n",
    "        obs_shape = n_in\n",
    "        current_obs = torch.zeros(\n",
    "            NUM_PROCESSES, obs_shape)  # torch.Size([16, 4])\n",
    "        rollouts = RolloutStorage(\n",
    "            NUM_ADVANCED_STEP, NUM_PROCESSES, obs_shape)  # rolloutsのオブジェクト\n",
    "        episode_rewards = torch.zeros([NUM_PROCESSES, 1])  # 現在の試行の報酬を保持\n",
    "        final_rewards = torch.zeros([NUM_PROCESSES, 1])  # 最後の試行の報酬を保持\n",
    "        obs_np = np.zeros([NUM_PROCESSES, obs_shape])  # Numpy配列\n",
    "        reward_np = np.zeros([NUM_PROCESSES, 1])  # Numpy配列\n",
    "        done_np = np.zeros([NUM_PROCESSES, 1])  # Numpy配列\n",
    "        each_step = np.zeros(NUM_PROCESSES)  # 各環境のstep数を記録\n",
    "        episode = 0  # 環境0の試行数\n",
    "\n",
    "        # 初期状態の開始\n",
    "        obs = [envs[i].reset() for i in range(NUM_PROCESSES)]\n",
    "        obs = np.array(obs)\n",
    "        obs = torch.from_numpy(obs).float()  # torch.Size([16, 4])\n",
    "        current_obs = obs  # 最新のobsを格納\n",
    "\n",
    "        # advanced学習用のオブジェクトrolloutsの状態の1つ目に、現在の状態を保存\n",
    "        rollouts.observations[0].copy_(current_obs)\n",
    "\n",
    "        # 実行ループ\n",
    "        for j in range(NUM_EPISODES*NUM_PROCESSES):  # 全体のforループ\n",
    "            # advanced学習するstep数ごとに計算\n",
    "            for step in range(NUM_ADVANCED_STEP):\n",
    "\n",
    "                # 行動を求める\n",
    "                with torch.no_grad():\n",
    "                    action = actor_critic.act(rollouts.observations[step])\n",
    "\n",
    "                # (16,1)→(16,)→tensorをNumPyに\n",
    "                actions = action.squeeze(1).numpy()\n",
    "\n",
    "                # 1stepの実行\n",
    "                for i in range(NUM_PROCESSES):\n",
    "                    obs_np[i], reward_np[i], done_np[i], _ = envs[i].step(\n",
    "                        actions[i])\n",
    "\n",
    "                    # episodeの終了評価と、state_nextを設定\n",
    "                    if done_np[i]:  # ステップ数が200経過するか、一定角度以上傾くとdoneはtrueになる\n",
    "\n",
    "                        # 環境0のときのみ出力\n",
    "                        if i == 0:\n",
    "                            print('%d Episode: Finished after %d steps' % (\n",
    "                                episode, each_step[i]+1))\n",
    "                            episode += 1\n",
    "\n",
    "                        # 報酬の設定\n",
    "                        if each_step[i] < 195:\n",
    "                            reward_np[i] = -1.0  # 途中でこけたら罰則として報酬-1を与える\n",
    "                        else:\n",
    "                            reward_np[i] = 1.0  # 立ったまま終了時は報酬1を与える\n",
    "\n",
    "                        each_step[i] = 0  # step数のリセット\n",
    "                        obs_np[i] = envs[i].reset()  # 実行環境のリセット\n",
    "\n",
    "                    else:\n",
    "                        reward_np[i] = 0.0  # 普段は報酬0\n",
    "                        each_step[i] += 1\n",
    "\n",
    "                # 報酬をtensorに変換し、試行の総報酬に足す\n",
    "                reward = torch.from_numpy(reward_np).float()\n",
    "                episode_rewards += reward\n",
    "\n",
    "                # 各実行環境それぞれについて、doneならmaskは0に、継続中ならmaskは1にする\n",
    "                masks = torch.FloatTensor(\n",
    "                    [[0.0] if done_ else [1.0] for done_ in done_np])\n",
    "\n",
    "                # 最後の試行の総報酬を更新する\n",
    "                final_rewards *= masks  # 継続中の場合は1をかけ算してそのまま、done時には0を掛けてリセット\n",
    "                # 継続中は0を足す、done時にはepisode_rewardsを足す\n",
    "                final_rewards += (1 - masks) * episode_rewards\n",
    "\n",
    "                # 試行の総報酬を更新する\n",
    "                episode_rewards *= masks  # 継続中のmaskは1なのでそのまま、doneの場合は0に\n",
    "\n",
    "                # 現在の状態をdone時には全部0にする\n",
    "                current_obs *= masks\n",
    "\n",
    "                # current_obsを更新\n",
    "                obs = torch.from_numpy(obs_np).float()  # torch.Size([16, 4])\n",
    "                current_obs = obs  # 最新のobsを格納\n",
    "\n",
    "                # メモリオブジェクトに今stepのtransitionを挿入\n",
    "                rollouts.insert(current_obs, action.data, reward, masks)\n",
    "\n",
    "            # advancedのfor loop終了\n",
    "\n",
    "            # advancedした最終stepの状態から予想する状態価値を計算\n",
    "\n",
    "            with torch.no_grad():\n",
    "                next_value = actor_critic.get_value(\n",
    "                    rollouts.observations[-1]).detach()\n",
    "                # rollouts.observationsのサイズはtorch.Size([6, 16, 4])\n",
    "\n",
    "            # 全stepの割引報酬和を計算して、rolloutsの変数returnsを更新\n",
    "            rollouts.compute_returns(next_value)\n",
    "\n",
    "            # ネットワークとrolloutの更新\n",
    "            global_brain.update(rollouts)\n",
    "            rollouts.after_update()\n",
    "\n",
    "            # 全部のNUM_PROCESSESが200step経ち続けたら成功\n",
    "            if final_rewards.sum().numpy() >= NUM_PROCESSES:\n",
    "                print('連続成功')\n",
    "                break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "\u001b[33mWARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.\u001b[0m\n",
      "0 Episode: Finished after 11 steps\n",
      "1 Episode: Finished after 11 steps\n",
      "2 Episode: Finished after 12 steps\n",
      "3 Episode: Finished after 41 steps\n",
      "4 Episode: Finished after 23 steps\n",
      "5 Episode: Finished after 15 steps\n",
      "6 Episode: Finished after 11 steps\n",
      "7 Episode: Finished after 10 steps\n",
      "8 Episode: Finished after 16 steps\n",
      "9 Episode: Finished after 15 steps\n",
      "10 Episode: Finished after 24 steps\n",
      "11 Episode: Finished after 30 steps\n",
      "12 Episode: Finished after 23 steps\n",
      "13 Episode: Finished after 32 steps\n",
      "14 Episode: Finished after 27 steps\n",
      "15 Episode: Finished after 66 steps\n",
      "16 Episode: Finished after 29 steps\n",
      "17 Episode: Finished after 16 steps\n",
      "18 Episode: Finished after 20 steps\n",
      "19 Episode: Finished after 17 steps\n",
      "20 Episode: Finished after 28 steps\n",
      "21 Episode: Finished after 19 steps\n",
      "22 Episode: Finished after 98 steps\n",
      "23 Episode: Finished after 58 steps\n",
      "24 Episode: Finished after 46 steps\n",
      "25 Episode: Finished after 132 steps\n",
      "26 Episode: Finished after 99 steps\n",
      "27 Episode: Finished after 200 steps\n",
      "28 Episode: Finished after 16 steps\n",
      "29 Episode: Finished after 68 steps\n",
      "30 Episode: Finished after 17 steps\n",
      "31 Episode: Finished after 200 steps\n",
      "32 Episode: Finished after 200 steps\n",
      "連続成功\n"
     ]
    }
   ],
   "source": [
    "# main学習\n",
    "cartpole_env = Environment()\n",
    "cartpole_env.run()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
